-module(quardtree).

-include("quardtree.hrl").

-include_lib("eunit/include/eunit.hrl").

-export([depth/1, sum/1, all_nodes/1, nth_node/2, leaves/1, nlevel_nodes/2, get_children/2, sub_quardtree/2]).

-spec depth(quardtree()) -> depth_type().
depth(#quardtree{depth=_Depth, nodes=Nodes}) ->
    trunc(math:log(3 * length(Nodes) + 1) / math:log(4)).

-spec sum(quardtree()) -> non_neg_integer().
sum(#quardtree{depth=Depth, nodes=Nodes}) ->
    trunc((math:pow(4, Depth) - 1) / 3).

-spec all_nodes(quardtree()) -> nodes_type().
all_nodes(#quardtree{depth=_Depth, nodes=Nodes}) ->
    Nodes.

-spec nth_node(level_type(), quardtree()) -> node_type().
nth_node(N, #quardtree{depth=_Depth, nodes=Nodes}) ->
    lists:nth(N, Nodes).

get_nodes(IndexList, Qt) ->
    get_nodes(IndexList, Qt, []).

get_nodes([], _Qt, Res) ->
    Res;
get_nodes([H|T], Qt, Res) ->
    get_nodes(T, Qt, Res ++ [nth_node(H, Qt)]).

-spec leaves(quardtree()) -> nodes_type().
leaves(#quardtree{depth=Depth,nodes=Nodes}) ->
    Start = trunc((math:pow(4, Depth-1) - 1) / 3),
    lists:nthtail(Start, Nodes).

-spec nlevel_nodes(pos_integer(), quardtree()) -> nodes_type().
nlevel_nodes(Level, #quardtree{depth=_Depth,nodes=Nodes}) ->
    Start = trunc((math:pow(4, Level-1) - 1) / 3) + 1,
    lists:sublist(Nodes, Start, trunc(math:pow(4, Level-1))).

-spec get_children(pos_integer(), quardtree()) -> nodes_type().
get_children(Nth, #quardtree{depth=Depth,nodes=Nodes}) ->
    case Nth =< (trunc(math:pow(4,Depth)-1)/3) of
        true ->
            Start = (Nth - 1) * 4 + 2,
            lists:sublist(Nodes, Start, 4);
        false ->
            []
    end.

get_nthnode_index_children(#quardtree{depth=Depth,nodes=_Nodes}, Nth) ->
    case Nth =< (trunc(math:pow(4,Depth-1)-1)/3) of
        true ->
            Start = (Nth - 1) * 4 + 2,
            [Start, Start+1, Start+2, Start+3];
        false ->
            []
    end.

-spec sub_quardtree(pos_integer(), quardtree()) -> quardtree().
sub_quardtree(NthNode, Qt) ->
    Nodes = get_descendants(Qt, [], [NthNode]),
    Depth = trunc(math:log(3 * length(Nodes) + 1) / math:log(4)),
    #quardtree{depth=Depth, nodes=Nodes}.
    
get_descendants(_Qt, Result, []) ->
    Result;
get_descendants(Qt, Result, IndexList) ->
    get_descendants(Qt, Result ++ get_nodes(IndexList,Qt), get_children_index(Qt, IndexList)).

get_children_index(Qt, L) ->
    get_children_index(Qt, [], L).

get_children_index(_Qt, Result, []) ->
    Result;
get_children_index(Qt, Result, [H|T]) ->
    get_children_index(Qt, Result ++ get_nthnode_index_children(Qt, H), T).



children_test() ->
    ?assertEqual([2,3,4,5], get_children(1, #quardtree{depth=1,nodes=[1,2,3,4,5]})),
    ?assertEqual([], get_children(2, #quardtree{depth=1,nodes=[47]})),
    ?assertEqual([],
        get_children(1, #quardtree{depth=1,nodes=[47]})),
    ?assertEqual([52,53,54,55],
        get_children(6, #quardtree{depth=4,nodes=lists:seq(31,105)})).

nodes_test() ->
    Qt = #quardtree{depth=4,nodes=lists:seq(31,115)},
    ?assertEqual(lists:seq(52,115),
        get_nodes(lists:seq(22,85), Qt)),
    ?assertEqual([31,32,33],
        get_nodes([1,2,3], Qt)).

descendants_test() ->
    ?assertEqual(#quardtree{depth=2,nodes=[1,2,3,4,5]}, sub_quardtree(1, #quardtree{depth=2,nodes=[1,2,3,4,5]})),
    ?assertEqual(#quardtree{depth=2,nodes=[36,52,53,54,55]},
        sub_quardtree(6, #quardtree{depth=4,nodes=lists:seq(31,115)})),
    ?assertEqual(#quardtree{depth=2,nodes=[37,56,57,58,59]},
        sub_quardtree(7, #quardtree{depth=4,nodes=lists:seq(31,115)})),
    ?assertEqual(#quardtree{depth=3,nodes=[32,36,37,38,39] ++ lists:seq(52,67)},
        sub_quardtree(2, #quardtree{depth=4,nodes=lists:seq(31,115)})).

chinldren_index_test() ->
    ?assertEqual([2,3,4,5], get_nthnode_index_children(#quardtree{depth=2,nodes=[11,22,33,44,55]}, 1)),
    ?assertEqual([], get_nthnode_index_children(#quardtree{depth=2,nodes=[11,22,33,44,55]}, 2)),
    ?assertEqual([2,3,4,5], get_children_index(#quardtree{depth=2,nodes=[11,22,33,44,55]},[1])),
    ?assertEqual([], get_children_index(#quardtree{depth=2,nodes=[11,22,33,44,55]},[2,3,4,5])).

leaves_test() ->
    ?assertEqual([1,2,3,4], 
        leaves(#quardtree{depth=2,nodes=[0,1,2,3,4]})),
    ?assertEqual([5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20], 
        leaves(#quardtree{depth=3,nodes=lists:seq(0,20)})).

level_nodes_test() ->
    ?assertEqual([0], 
        nlevel_nodes(1,#quardtree{depth=2,nodes=[0,1,2,3,4]})),
    ?assertEqual([1,2,3,4],
        nlevel_nodes(2, #quardtree{depth=2,nodes=[0,1,2,3,4]})),
    ?assertEqual(lists:seq(5,20),
        nlevel_nodes(3, #quardtree{depth=3, nodes=lists:seq(0,20)})).

nodes_leaves_test() ->
    Qt = #quardtree{depth=1, nodes=[0]},
    ?assertEqual(leaves(Qt), 
        nlevel_nodes(Qt#quardtree.depth, Qt)),
    Qt2 = #quardtree{depth=2, nodes=[0,1,2,3,4]},
    ?assertEqual(leaves(Qt2), 
        nlevel_nodes(Qt2#quardtree.depth, Qt2)),
    Qt3 = #quardtree{depth=3, nodes=lists:seq(0,20)},
    ?assertEqual(leaves(Qt3), 
        nlevel_nodes(Qt3#quardtree.depth, Qt3)).

nthnode_leaves_test() ->
    Qt = #quardtree{depth=1, nodes=[0]},
    ?assertEqual([], get_children(1, Qt)),
    Qt3 = #quardtree{depth=3, nodes=lists:seq(0,20)},
    ?assertEqual([1,2,3,4], get_children(1, Qt3)),
    ?assertEqual([9,10,11,12], get_children(3, Qt3)),
    ?assertEqual([13,14,15,16], get_children(4, Qt3)),
    ?assertEqual([17,18,19,20], get_children(5, Qt3)).

